import streamlit as st
import time
import plotly.express as px
from basic_predict_fun import predict
from utils.config import Config

# 页面配置
st.set_page_config(page_title="NLP 分类系统", layout="wide")

st.title("🤖 NLP 分类过程演示")
st.markdown("支持 **cat 多分类 (随机森林)** 和 **label 二分类 (逻辑回归)** 的预测。")

# 初始化 session_state
if "result" not in st.session_state:
    st.session_state.result = None
if "show_final" not in st.session_state:
    st.session_state.show_final = False

# 输入文本
text = st.text_area("请输入一段文本进行预测", height=200, key="input_text")
task = st.radio("请选择任务", ["cat", "label", "all"], index=2, key="task_select")

root_path = 'D:/pycode/group4_nlp_project'
conf = Config(root_path)
thinking_img = conf.thinking_img  # 思考中
final_img = conf.final_img        # 最终判断

# 点击按钮后执行预测
if st.button("开始预测"):
    if text.strip() == "":
        st.warning("请输入文本！")
    else:
        with st.spinner("🧠 正在思考中..."):
            st.image(thinking_img, width=300)
            time.sleep(2)  # 模拟推理时间
            st.session_state.result = predict(text, task=task)
            st.session_state.show_final = False

# 如果有预测结果，展示
if st.session_state.result:
    result = st.session_state.result
    st.markdown("---")
    st.subheader("🎯 预测结果")

    # --- 多分类结果 ---
    if "cat" in result:
        if "error" in result["cat"]:
            st.error(result["cat"]["error"])
        else:
            best = result["cat"]["best_category"]
            best_prob = result["cat"]["best_probability"]
            st.success(f"预测类别: {best} (置信度 {best_prob:.4f})")

            # 展示 Top-3
            top3 = result["cat"]["top3"]
            categories = [item["category"] for item in top3]
            probs = [item["probability"] for item in top3]

            col1, col2 = st.columns(2)

            with col1:
                st.write("📊 Top-3 条形图")
                fig_bar = px.bar(
                    x=categories, y=probs,
                    text=[f"{p:.2%}" for p in probs],
                    labels={"x": "类别", "y": "概率"},
                    color=categories,
                    color_discrete_sequence=px.colors.qualitative.Set2  # ✅ 更清晰的调色板
                )
                fig_bar.update_traces(textposition="outside")
                st.plotly_chart(fig_bar, use_container_width=True)

            with col2:
                st.write("🥧 Top-3 饼图")
                fig_pie = px.pie(
                    names=categories, values=probs,
                    color=categories,
                    color_discrete_sequence=px.colors.qualitative.Set2,  # ✅ 更清晰的调色板
                    hole=0.3
                )
                fig_pie.update_traces(textinfo="label+percent")  # 显示标签+百分比
                st.plotly_chart(fig_pie, use_container_width=True)

    # --- 二分类结果 ---
    if "label" in result:
        if "error" in result["label"]:
            st.error(result["label"]["error"])
        else:
            label = result["label"]["label"]
            prob = result["label"]["probability"]
            st.success(f"预测标签: {label} (置信度 {prob:.4f})")

    # --- JSON 展示 ---
    with st.expander("查看完整结果 JSON"):
        st.json(result)

    # 延迟显示最终结论
    if not st.session_state.show_final:
        time.sleep(1.5)
        st.session_state.show_final = True

    if st.session_state.show_final:
        st.markdown("---")
        st.subheader("✅ 我的最终判断")
        st.image(final_img, width=300)

        # --- 多分类最终结果 ---
        if "cat" in result and "best_category" in result["cat"]:
            best = result["cat"]["best_category"]
            prob = result["cat"]["best_probability"]
            st.markdown(
                f"<h2 style='color:#4CAF50;'>📚 我认为这是 <b>{best}</b> ，其概率为 ({prob:.4f})</h2>",
                unsafe_allow_html=True
            )

        # --- 二分类最终结果 ---
        if "label" in result:
            label = result["label"]["label"]
            prob = result["label"]["probability"]
            st.markdown(
                f"<h2 style='color:#2196F3;'>💬 我认为这是 <b>{label}</b> ，其概率为 ({prob:.4f})</h2>",
                unsafe_allow_html=True
            )
